import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import geopandas as gpd
import shapely.geometry as sgeom

"""
This script is to describe the age and geographical distribution of depression patients in China.
created by JINJing
"""

# 生成模拟数据
np.random.seed(0)
n_patients = 1000  # 假设有1000名患者
ages = np.random.randint(10, 81, size=n_patients)  # 年龄范围从10到80岁
provinces = np.random.choice(['北京', '上海', '广东', '江苏', '浙江', '四川', '河南', '山东', '湖北', '湖南'],
                             size=n_patients)  # 假设这些省份有抑郁症患者

# 将数据整合到DataFrame中
data = pd.DataFrame({
    'age': ages,
    'province': provinces
})

# 计算年龄分布
age_bins = np.arange(10, 91, 10)  # 定义年龄区间
age_distribution = data['age'].value_counts(bins=age_bins, sort=False)
age_distribution.index = age_bins[:-1] + 5  # 设置区间的中点为索引
age_distribution = age_distribution.sort_index()  # 按年龄排序

# 绘制年龄分布图
plt.figure(figsize=(10, 6))
age_distribution.plot(kind='bar', color='skyblue')
plt.xlabel('Age')
plt.ylabel('Number of Patients')
plt.title('Age Distribution of Depression Patients')
plt.show()

# 使用模拟数据创建简单的GeoDataFrame来表示省份
provinces_geometry = {
    'province': ['北京', '上海', '广东', '江苏', '浙江', '四川', '河南', '山东', '湖北', '湖南'],
    'geometry': [sgeom.Point(random_lon, random_lat) for random_lon, random_lat in
                 zip(np.random.rand(10) * 180 - 90, np.random.rand(10) * 360 - 180)]
}
gdf = gpd.GeoDataFrame(provinces_geometry, geometry='geometry')

# 将模拟数据合并到GeoDataFrame中（这里使用患者数量的计数）
gdf['patient_count'] = data['province'].value_counts()

# 绘制全国分布图（使用点来表示省份）
fig, ax = plt.subplots(1, 1, figsize=(10, 6))
gdf.plot(ax=ax, column='patient_count', cmap='YlOrRd', markersize=50, linewidth=0.8, edgecolor='0.8')
ax.set_title('National Distribution of Depression Patients (Simulated)')
plt.show()